package com.heima.leetcode.datastructure.heap;

import java.util.Iterator;

/**
 * 堆，可以根据情况创建小顶堆或者大顶堆
 * @author 勾新杰
 * @version 1.0
 * @date 2024/9/6 21:53
 */
public class Heap implements Iterable<Integer> {
    private int[] arr;
    private int size;
    private final boolean max; // 时大顶堆还是小顶堆
    public Heap(int capacity, boolean max) {
        arr = new int[capacity];
        this.max = max;
    }

    public Heap(int[] arr, boolean max) {
        this.arr = arr;
        size = arr.length;
        this.max = max;
        heapify();
    }

    public boolean isFull(){
        return size == arr.length;
    }

    public boolean isEmpty(){
        return size == 0;
    }

    public void offer(int value){
        if (isFull()){
            // 扩容
            grow();
        }
        up(size++,value);
    }

    public int poll(){
        if (isEmpty()){
            throw new IllegalArgumentException("堆为空");
        }
        int value = arr[0];
        swap(0,--size);
        down(0);
        return value;
    }

    public int peek(){
        if (isEmpty()){
            throw new IllegalArgumentException("堆为空");
        }
        return arr[0];
    }

    public int size(){
        return size;
    }

    private void heapify(){
        for (int i = size / 2 - 1; i >= 0; i--){
            down(i);
        }
    }

    private void down(int parent){
        int leftChild = parent * 2 + 1;
        int rightChild = leftChild + 1;
        int maxOrMin = max ? maxValue(parent,leftChild,rightChild) : minValue(parent,leftChild,rightChild);
        if (maxOrMin == parent){
            return;
        }
        swap(parent,maxOrMin);
        down(maxOrMin);
    }

    private void up(int index,int value){
        int child = index;
        int parent = (child - 1) / 2;
        while (child > 0 && (max ? arr[parent] < value : arr[parent] > value)){
            swap(parent,child);
            child = parent;
            parent = (child - 1) / 2;
        }
        arr[child] = value;
    }

    private int maxValue(int parent, int leftChild, int rightChild){
        if (leftChild >= size){
            return parent;
        }
        if (rightChild == size){
            return arr[parent] >= arr[leftChild] ? parent : leftChild; // 相等减少交换
        }
        if (arr[parent] >= arr[leftChild]){
            return arr[parent] >= arr[rightChild] ? parent : rightChild; // 相等减少交换
        }
        return arr[leftChild] >= arr[rightChild] ? leftChild : rightChild; // 相等走左边分支
    }

    private int minValue(int parent, int leftChild, int rightChild){
        if (leftChild >= size){
            return parent;
        }
        if (rightChild == size){
            return arr[parent] <= arr[leftChild] ? parent : leftChild; // 相等减少交换
        }
        if (arr[parent] <= arr[leftChild]){
            return arr[parent] <= arr[rightChild] ? parent : rightChild; // 相等减少交换
        }
        return arr[leftChild] <= arr[rightChild] ? leftChild : rightChild; // 相等走左边分支
    }

    private void swap(int index1, int index2){
        int temp = arr[index1];
        arr[index1] = arr[index2];
        arr[index2] = temp;
    }

    private void grow(){ // 默认扩容1.5倍
         int capacity = arr.length + (arr.length >> 1);
         int[] array = new int[capacity];
         System.arraycopy(arr,0,array,0,arr.length);
         arr = array;
    }

    @Override
    public Iterator<Integer> iterator() {
        return new Iterator<Integer>() {
            int index = 0;
            @Override
            public boolean hasNext() {
                return index != size;
            }

            @Override
            public Integer next() {
                return arr[index++];
            }
        };
    }

    @Override
    public String toString() {
        StringBuilder builder = new StringBuilder();
        builder.append("[");
        if (!max) {
            for (int i = 0; i < size; i++) {
                builder.append(arr[i]);
                if (i != size - 1){
                    builder.append(", ");
                }
            }
        }else {
            for (int i = size - 1; i >= 0; i--) {
                builder.append(arr[i]);
                if (i != 0){
                    builder.append(", ");
                }
            }
        }
        builder.append("]");
        return builder.toString();
    }
}
